# -*- coding: utf-8 -*-
"""
Created on Wed Jul 15 16:22:36 2020

@author: Administrator
"""
import avadata as avadata
import cv2
import numpy as np
from collections import defaultdict
from arg import cmd
from tqdm import tqdm
all_path=cmd()

path_label=all_path.label #r'C:/Users/Administrator/Downloads/haikang/ava/label/ava-action-categories.txt'
path_annotation=all_path.annotation
videoname= all_path.name    
endsec=all_path.endsec
all_box=avadata.load_box_by_name(path_annotation)
class videoReader(object):
    def __init__(self, source):
        self.source=source
        self.cap=cv2.VideoCapture(self.source)
        self.display_width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        self.display_height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        self.fps=round(self.cap.get(cv2.CAP_PROP_FPS))
        self.cap.set(cv2.CAP_PROP_POS_FRAMES,902*self.fps) #从902秒开始读
        self.endtime=1000
        if all_path.endsec:
            self.endtime=endsec
        self.output_file=None
        if all_path.outpath:
            self.output_file=self.get_output_file(r'C:/Users/Administrator/Downloads/haikang/ava/output.mp4')
        #r'C:/Users/Administrator/Downloads/haikang/ava/output.mp4'
        if not self.cap.isOpened():
            raise IOError("Video {} cannot be opened".format(self.source))
    def __iter__(self):
        return self
    
    def __next__(self):
        pos = self.cap.get(cv2.CAP_PROP_POS_FRAMES)#获取当前帧号
        was_read,frame =self.cap.read()
        
        if pos < self.endtime*self.fps:
            
            return was_read, frame, pos
        else:
            raise StopIteration
    
        if not was_read :
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
            frame = None
        return was_read, frame, pos
    def get_output_file(self, path):
        return cv2.VideoWriter(
            filename=path,
            fourcc=cv2.VideoWriter_fourcc(*"mp4v"),
            fps=float(30),
            frameSize=(self.display_width, self.display_height),
            isColor=True,
        )
        
    def display(self, frame):

        if self.output_file is None:
            cv2.imshow("SlowFast", frame)
        else:
            self.output_file.write(frame)
       
    def clean(self):
        self.cap.release()
        
        if self.output_file is None:
            cv2.destroyAllWindows()
        else:
            self.output_file.release()
        
def getitem_by_name(name, mode_on=True):# 运行文件后执行这个函数，注意要选一个训练集里有的视频
    '''
    arg:
        name of the video (e.g.'BCiuXAuCKAU')
        only label the keyframe if mode_on=True else label every frames according to the sec they belong to. 
    return:
        no return but will output a frame with box and label 
    '''
    name_ele=name.split('/')[-1].split('.')[0] #从路径中获取视频的名称
    
    assert name_ele in all_box.keys( )
    labellist=[]
    with open(path_label) as f:
        for line in f.readlines():
            labellist.append(line.strip('\n'))
    
    item=all_box[name_ele]
    
    frame_provider=videoReader(name)
    
    box_sec=defaultdict(list)
    label_sec=defaultdict(list)
    for sec in item.keys():
        if len(item[sec]) ==0:
            box_sec[sec]
            label_sec[sec]
        for ele in item[sec]:
            box_sec[sec].append(ele[:4])
            label_sec[sec].append(list(map(int,ele[4:-1])))
    for able_to_read, frame, pos in tqdm(frame_provider,desc='processing frame'): 
        
        if not able_to_read:
            frame=[]
            break
        index= int(pos // frame_provider.fps)
        sub_index=pos/frame_provider.fps
        if( mode_on and sub_index not in item.keys() )or len(item[index])==0: #or之前是为了只框选关键帧，or之后是为了处理没有被label的关键帧
            frame_provider.display(frame)
            continue
        else:
            
                
            boxes=box_sec[index]
            labelset=label_sec[index]
            
    
            boxes=np.asarray(boxes)
            boxes[:,[0,2]]*=frame_provider.display_width
            boxes[:,[1,3]]*=frame_provider.display_height
            boxes=clip_boxes_to_image(boxes, frame_provider.display_height, frame_provider.display_width)
            labelset=np.asarray(labelset)
            
            
            
            #print(index)
            #print(boxes)
            #print(labelset)
            for i,(box, box_labels) in enumerate(zip(boxes.astype(int), labelset)):
                cv2.rectangle(
                    frame,
                    tuple(box[:2]),
                    tuple(box[2:]),
                    (0,255,0),
                    thickness=2,
                )
                
                label_origin = box[:2]
                label_origin[-1]-= 10*i
                for label in box_labels:
                    labelname=labellist[label-1]
                    
                    (label_width, label_height), _ = cv2.getTextSize(
                        labelname, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
                    )
                    label_origin[-1] -=  5
                    '''
                    cv2.rectangle(
                        frame,
                        (label_origin[0], label_origin[1]+5),
                        (
                            label_origin[0]+label_width,
                            label_origin[1]-label_height-5,
                            
                        ),
                        (255,255,255),
                        -1,
                    )   
                    '''
                    cv2.putText(
                        frame,
                        labelname,
                        tuple(label_origin),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.5,
                        (0, 0, 255),
                        3,
                    )
                    label_origin[-1] -= label_height + 5
                    
        if mode_on:           
            for i in range(30):
                frame_provider.display(frame)
        else:
            frame_provider.display(frame)
        #key = cv2.waitKey(1)
        #if key == 27:
            #break
    frame_provider.clean()





def clip_boxes_to_image(boxes, height, width):

    """

    Clip the boxes with the height and width of the image size.

    Args:

        boxes (ndarray): bounding boxes to peform crop. The dimension is

        `num boxes` x 4.

        height (int): the height of the image.

        width (int): the width of the image.

    Returns:

        boxes (ndarray): cropped bounding boxes.

    """

    boxes[:, [0, 2]] = np.minimum(

        width - 1.0, np.maximum(0.0, boxes[:, [0, 2]])

    )

    boxes[:, [1, 3]] = np.minimum(

        height - 1.0, np.maximum(0.0, boxes[:, [1, 3]])

    )

    return boxes
